import hashlib
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64
from typing import List, Dict, Union


class CustomAlgorithm:
    def __init__(self, h5st_algo_config: Dict):
        self.h5st_algo_config = h5st_algo_config

    def add_salt(self, message: Union[str, bytes]) -> Union[str, bytes]:
        salt = self.h5st_algo_config.get('customAlgorithm', {}).get('salt', '')
        if isinstance(message, str):
            return message + salt
        return message

    def e_key(self, key: str) -> str:
        convert_index = self.h5st_algo_config.get('customAlgorithm', {}).get('convertIndex', {}).get('hmac', None)
        if not convert_index:
            return key

        split = list(key)
        slice1 = split[:convert_index]
        slice2 = split[convert_index:]
        array = [chr(158 - ord(char)) for char in slice1[::-1]]

        return ''.join(array + slice2)

    def parse_loop(self, base64_str: str, base64_str_length: int, reverse_map: Dict[int, int]):
        words = []
        n_bytes = 0
        for i in range(base64_str_length):
            if i % 4:
                bits1 = reverse_map[ord(base64_str[i - 1])] << ((i % 4) * 2)
                bits2 = reverse_map[ord(base64_str[i])] >> (6 - (i % 4) * 2)
                bits_combined = bits1 | bits2
                if len(words) <= n_bytes >> 2:
                    words.append(0)
                words[n_bytes >> 2] |= bits_combined << (24 - (n_bytes % 4) * 8)
                n_bytes += 1
        return words[:n_bytes]

    def aes_encrypt(self, message: Union[str, bytes], key: str, iv: bytes = b'\0' * 16) -> str:
        if isinstance(message, str):
            message = message.encode()
        key = key[::-1] if self.h5st_algo_config.get('customAlgorithm', {}).get('keyReverse') else key
        cipher = AES.new(key.encode(), AES.MODE_CBC, iv)
        encrypted = cipher.encrypt(pad(message, AES.block_size))
        return base64.b64encode(encrypted).decode()

    def aes_decrypt(self, ciphertext: str, key: str, iv: bytes = b'\0' * 16) -> str:
        ciphertext = base64.b64decode(ciphertext)
        key = key[::-1] if self.h5st_algo_config.get('customAlgorithm', {}).get('keyReverse') else key
        cipher = AES.new(key.encode(), AES.MODE_CBC, iv)
        decrypted = unpad(cipher.decrypt(ciphertext), AES.block_size)
        return decrypted.decode()

    def md5(self, message: Union[str, bytes]) -> str:
        return hashlib.md5(self.add_salt(message).encode()).hexdigest()

    def sha256(self, message: Union[str, bytes]) -> str:
        return hashlib.sha256(self.add_salt(message).encode()).hexdigest()

    def hmac_sha256(self, message: Union[str, bytes], key: str) -> str:
        key = self.e_key(key)
        return hashlib.sha256(self.add_salt(message).encode()).hexdigest()

    def encode_base64(self, word_array: List[int]) -> str:
        map_chars = self.h5st_algo_config.get('customAlgorithm', {}).get('map', '')
        if not map_chars:
            raise ValueError("This version of the algorithm lacks required parameters.")

        byte_array = bytearray(word_array)
        base64_chars = []
        for i in range(0, len(byte_array), 3):
            triplet = byte_array[i:i + 3]
            if len(triplet) < 3:
                triplet.extend([0] * (3 - len(triplet)))
            enc_val = int.from_bytes(triplet, 'big')
            for j in range(4):
                base64_chars.append(map_chars[(enc_val >> (18 - 6 * j)) & 0x3f])

        result = []
        for i in range(0, len(base64_chars), 4):
            result.extend(reversed(base64_chars[i:i + 4]))

        return ''.join(result)

    def decode_base64(self, base64_str: str) -> List[int]:
        map_chars = self.h5st_algo_config.get('customAlgorithm', {}).get('map', '')
        if not map_chars:
            raise ValueError("This version of the algorithm lacks required parameters.")

        reverse_map = {ord(c): i for i, c in enumerate(map_chars)}
        base64_str = ''.join(reversed([base64_str[i:i + 4] for i in range(0, len(base64_str), 4)]))
        padding_char = map_chars[64]
        if padding_char in base64_str:
            base64_str = base64_str[:base64_str.index(padding_char)]

        word_array = self.parse_loop(base64_str, len(base64_str), reverse_map)
        return word_array


